﻿# coding=utf-8
# 2023/4/27: 首个记录，基于旧2020/8/5版本
# 2023/5/5: 枚举基于Enum
# 2023/10/10: 补充部分英文注释
# 2025/7/21: 修正orientation插值

import bi_common as bi
from math import pi, sin, cos, atan2
from enum import Enum


class DynamicObjectType(Enum):
    UNKNOWN = 0  # 未知
    CAR = 1  # 轿车
    VAN = 2  # 面包车
    TRUCK = 3  # 货车
    TRAILER = 4  # 拖车
    SEMITRAILER = 5  # 货柜车
    BUS = 6  # 大巴
    MOTORBIKE = 7  # 摩托车
    BICYCLE = 8  # 自行车
    TRAIN = 9  # 火车
    TRAM = 10  # 有轨电车
    PEDESTRIAN = 11  # 行人
    WHEELCHAIR = 12  # 轮椅
    ANIMAL = 13  # 动物


class DynamicObject:
    def __init__(self):
        self.id = 0  # ID of the object, 0 as subject, 1~ as other objects / 目标物ID，0表示为主目标物，1~表示一般目标物
        self.type = DynamicObjectType.UNKNOWN  # Classification of the object / 目标物类型
        self.center_longitude = 0.0  # [°E] Longitude of the object's box center / 目标物中心的经度
        self.center_latitude = 0.0  # [°N] Latitude of the object's box center / 目标物中心的纬度，北为正方向
        self.center_altitude_valid = False
        self.center_altitude = 0.0  # [m] Altitude of the object's box center / 目标物中心的海拔
        self.orientation_valid = False
        self.orientation = 0.0  # [deg] Heading of the object's box, 0 as north, CCW is positive / 目标物朝向，北为0，逆时针为正方向
        self.width_valid = False
        self.width = 0.0  # [m] Width of the object's box / 目标物宽度
        self.length_valid = False
        self.length = 0.0  # [m] Length of the object's box / 目标物长度
        self.height_valid = False
        self.height = 0.0  # [m] Height of the object's box / 目标物高度
        self.speed_valid = False
        self.speed = 0.0  # [KPH] Speed of the object's box along the heading / 目标物速度
        self.accel_x_valid = False
        self.accel_x = 0.0  # [m/s²] Acceleration of the object's box along the heading / 目标物纵向加速度
        self.curvature_valid = False
        self.curvature = 0.0  # [1/m] Turn curvature of the object's rear axis center / 目标物转弯曲率
        self.yaw_rate_valid = False
        self.yaw_rate = 0.0  # [deg/s] Yaw rate of the object's box / 目标物横摆角速度
        self.pitch_valid = False
        self.pitch = 0.0  # [deg] Pitch angle of the object's box, heading down is positive / 目标物俯仰角，车头朝下为正
        self.roll_valid = False
        self.roll = 0.0  # [deg] Roll angle of the object's box, right side heading down is positive / 目标物横滚角，右侧朝下为正
        self.rear_axis_longitudinal_offset_valid = False
        self.rear_axis_longitudinal_offset = 0.0  # [m] Longitudinal distance between the object's box center and rear axis center, always positive / 目标物自身坐标系下，后轴中心与物体中心的纵向距离，恒为正
        self.rear_axis_height_offset_valid = False
        self.rear_axis_height_offset = 0.0  # [m] Height difference of the object's box center and rear axis center, always positive / 目标物自身坐标系下，后轴中心与物体中心的高度差，恒为正
        self.color_valid = False
        self.color_r = 0  # Red component of the n-th object's color / 目标物颜色R分量
        self.color_g = 0  # Green component of the n-th object's color / 目标物颜色G分量
        self.color_b = 0  # Blue component of the n-th object's color / 目标物颜色B分量


class GlobalObjectsSample:
    def __init__(self):
        self.time = 0.0  # [s] Time offset in session / Session内的相对时间
        self.objects = []  # List of objects / 目标物列表

    # Convert to general sample for output / 转通用样本，用于样本输出
    def to_general_sample(self):
        output = bi.agency.create_general_sample()
        output.protocol = 'global-objects-sample-v2'
        output.time = self.time
        output.significant = 1 + len(self.objects) * 26
        output.values = []
        output.values.append(len(self.objects))
        for obj in self.objects:
            output.values.append(obj.id)
            output.values.append(obj.type.value)
            output.values.append(obj.center_longitude)
            output.values.append(obj.center_latitude)
            output.values.append(obj.center_altitude if obj.center_altitude_valid else None)
            output.values.append(obj.orientation if obj.orientation_valid else None)
            output.values.append(obj.width if obj.width_valid else None)
            output.values.append(obj.length if obj.length_valid else None)
            output.values.append(obj.height if obj.height_valid else None)
            output.values.append(obj.speed if obj.speed_valid else None)
            output.values.append(obj.accel_x if obj.accel_x_valid else None)
            output.values.append(obj.curvature if obj.curvature_valid else None)
            output.values.append(obj.yaw_rate if obj.yaw_rate_valid else None)
            output.values.append(obj.pitch if obj.pitch_valid else None)
            output.values.append(obj.roll if obj.roll_valid else None)
            output.values.append(obj.rear_axis_longitudinal_offset if obj.rear_axis_longitudinal_offset_valid else None)
            output.values.append(obj.rear_axis_height_offset if obj.rear_axis_height_offset_valid else None)
            output.values.append(obj.color_r if obj.color_valid else None)
            output.values.append(obj.color_g if obj.color_valid else None)
            output.values.append(obj.color_b if obj.color_valid else None)
            output.values.append(None)
            output.values.append(None)
            output.values.append(None)
            output.values.append(None)
            output.values.append(None)
            output.values.append(None)
        return output


def _conv_global_objects_sample_v1(gs):
    values_count = len(gs.values)
    if values_count < 1:
        return None
    object_count = int(gs.values[0]) if gs.values[0] is not None else 0
    if values_count != 1 + 10 * object_count:
        return None
    output = GlobalObjectsSample()
    output.time = gs.time
    for i in range(0, object_count):
        obj = DynamicObject()
        b = 1 + 10 * i
        obj.id = int(gs.values[b] if gs.values[b] is not None else 0)
        obj.type = DynamicObjectType(int(gs.values[b + 1])) if gs.values[b + 1] is not None else DynamicObjectType.UNKNOWN
        obj.center_longitude = float(gs.values[b + 2]) if gs.values[b + 2] is not None else 0.0
        obj.center_latitude = float(gs.values[b + 3]) if gs.values[b + 3] is not None else 0.0
        if gs.values[b + 4] is not None:
            obj.center_altitude_valid = True
            obj.center_altitude = float(gs.values[b + 4])
        if gs.values[b + 5] is not None:
            obj.orientation_valid = True
            obj.orientation = float(gs.values[b + 5])
        if gs.values[b + 6] is not None:
            obj.width_valid = True
            obj.width = float(gs.values[b + 6])
        if gs.values[b + 7] is not None:
            obj.length_valid = True
            obj.length = float(gs.values[b + 7])
        if gs.values[b + 8] is not None:
            obj.height_valid = True
            obj.height = float(gs.values[b + 8])
        if gs.values[b + 9] is not None:
            obj.speed_valid = True
            obj.speed = float(gs.values[b + 9])
        output.objects.append(obj)
    return output


def _conv_global_objects_sample_v2(gs):
    values_count = len(gs.values)
    if values_count < 1:
        return None
    object_count = int(gs.values[0]) if gs.values[0] is not None else 0
    if values_count != 1 + 26 * object_count:
        return None
    output = GlobalObjectsSample()
    output.time = gs.time
    for i in range(0, object_count):
        obj = DynamicObject()
        b = 1 + 26 * i
        obj.id = int(gs.values[b] if gs.values[b] is not None else 0)
        obj.type = DynamicObjectType(int(gs.values[b + 1])) if gs.values[b + 1] is not None else DynamicObjectType.UNKNOWN
        obj.center_longitude = float(gs.values[b + 2]) if gs.values[b + 2] is not None else 0.0
        obj.center_latitude = float(gs.values[b + 3]) if gs.values[b + 3] is not None else 0.0
        if gs.values[b + 4] is not None:
            obj.center_altitude_valid = True
            obj.center_altitude = float(gs.values[b + 4])
        if gs.values[b + 5] is not None:
            obj.orientation_valid = True
            obj.orientation = float(gs.values[b + 5])
        if gs.values[b + 6] is not None:
            obj.width_valid = True
            obj.width = float(gs.values[b + 6])
        if gs.values[b + 7] is not None:
            obj.length_valid = True
            obj.length = float(gs.values[b + 7])
        if gs.values[b + 8] is not None:
            obj.height_valid = True
            obj.height = float(gs.values[b + 8])
        if gs.values[b + 9] is not None:
            obj.speed_valid = True
            obj.speed = float(gs.values[b + 9])
        if gs.values[b + 10] is not None:
            obj.accel_x_valid = True
            obj.accel_x = float(gs.values[b + 10])
        if gs.values[b + 11] is not None:
            obj.curvature_valid = True
            obj.curvature = float(gs.values[b + 11])
        if gs.values[b + 12] is not None:
            obj.yaw_rate_valid = True
            obj.yaw_rate = float(gs.values[b + 12])
        if gs.values[b + 13] is not None:
            obj.pitch_valid = True
            obj.pitch = float(gs.values[b + 13])
        if gs.values[b + 14] is not None:
            obj.roll_valid = True
            obj.roll = float(gs.values[b + 14])
        if gs.values[b + 15] is not None:
            obj.rear_axis_longitudinal_offset_valid = True
            obj.rear_axis_longitudinal_offset = float(gs.values[b + 15])
        if gs.values[b + 16] is not None:
            obj.rear_axis_height_offset_valid = True
            obj.rear_axis_height_offset = float(gs.values[b + 16])
        if gs.values[b + 17] is not None and gs.values[b + 18] is not None and gs.values[b + 19] is not None:
            obj.color_valid = True
            obj.color_r = int(gs.values[b + 17])
            obj.color_g = int(gs.values[b + 18])
            obj.color_b = int(gs.values[b + 19])
        output.objects.append(obj)
    return output


def _interpolate_angle(a1, w1, a2, w2):
    deg2rad = pi / 180
    x1 = cos(a1 * deg2rad)
    y1 = sin(a1 * deg2rad)
    x2 = cos(a2 * deg2rad)
    y2 = sin(a2 * deg2rad)
    xo = x1 * w1 + x2 * w2
    yo = y1 * w1 + y2 * w2
    if xo == 0 and yo == 0:
        return None
    return atan2(yo, xo) / deg2rad


def _interpolate_global_objects_sample(s1, w1, s2, w2):
    output = GlobalObjectsSample()
    output.time = bi.time
    for o1 in s1.objects:
        id = o1.id
        for o2 in s2.objects:
            if o2.id != id:
                continue
            obj = DynamicObject()
            obj.id = id
            obj.type = o1.type if w1 > w2 else o2.type
            obj.center_longitude = o1.center_longitude * w1 + o2.center_longitude * w2
            obj.center_latitude = o1.center_latitude * w1 + o2.center_latitude * w2
            if o1.center_altitude_valid and o2.center_altitude_valid:
                obj.center_altitude_valid = True
                obj.center_altitude = o1.center_altitude * w1 + o2.center_altitude * w2
            if o1.orientation_valid and o2.orientation_valid:
                obj.orientation_valid = True
                obj.orientation = _interpolate_angle(o1.orientation, w1, o2.orientation, w2)
            if o1.width_valid and o2.width_valid:
                obj.width_valid = True
                obj.width = o1.width * w1 + o2.width * w2
            if o1.length_valid and o2.length_valid:
                obj.length_valid = True
                obj.length = o1.length * w1 + o2.length * w2
            if o1.height_valid and o2.height_valid:
                obj.height_valid = True
                obj.height = o1.height * w1 + o2.height * w2
            if o1.speed_valid and o2.speed_valid:
                obj.speed_valid = True
                obj.speed = o1.speed * w1 + o2.speed * w2
            if o1.accel_x_valid and o2.accel_x_valid:
                obj.accel_x_valid = True
                obj.accel_x = o1.accel_x * w1 + o2.accel_x * w2
            if o1.curvature_valid and o2.curvature_valid:
                obj.curvature_valid = True
                obj.curvature = o1.curvature * w1 + o2.curvature * w2
            if o1.yaw_rate_valid and o2.yaw_rate_valid:
                obj.yaw_rate_valid = True
                obj.yaw_rate = o1.yaw_rate * w1 + o2.yaw_rate * w2
            if o1.pitch_valid and o2.pitch_valid:
                obj.pitch_valid = True
                obj.pitch = o1.pitch * w1 + o2.pitch * w2
            if o1.roll_valid and o2.roll_valid:
                obj.roll_valid = True
                obj.roll = o1.roll * w1 + o2.roll * w2
            if o1.rear_axis_longitudinal_offset_valid and o2.rear_axis_longitudinal_offset_valid:
                obj.rear_axis_longitudinal_offset_valid = True
                obj.rear_axis_longitudinal_offset = o1.rear_axis_longitudinal_offset * w1 + o2.rear_axis_longitudinal_offset * w2
            if o1.rear_axis_height_offset_valid and o2.rear_axis_height_offset_valid:
                obj.rear_axis_height_offset_valid = True
                obj.rear_axis_height_offset = o1.rear_axis_height_offset * w1 + o2.rear_axis_height_offset * w2
            if o1.color_valid and o2.color_valid:
                obj.color_valid = True
                obj.color_r = o1.color_r if w1 > w2 else o2.color_r
                obj.color_g = o1.color_g if w1 > w2 else o2.color_g
                obj.color_b = o1.color_b if w1 > w2 else o2.color_b
            output.objects.append(obj)
    return output


# Query GlobalObjectsSample for input / 获取GlobalObjectsSample，用于样本输入
def get_global_objects_sample():
    s1 = None
    s2 = None
    w1 = 0.0
    w2 = 0.0
    protocol_id_v1 = 'global-objects-sample-v1'
    protocol_id = 'global-objects-sample-v2'
    if protocol_id in bi.input_samples:
        pair = bi.input_samples[protocol_id]
        s1 = _conv_global_objects_sample_v2(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_global_objects_sample_v2(pair.sample2)
        w2 = pair.weight2
    elif protocol_id_v1 in bi.input_samples:
        pair = bi.input_samples[protocol_id_v1]
        s1 = _conv_global_objects_sample_v1(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_global_objects_sample_v1(pair.sample2)
        w2 = pair.weight2
    if s1 is not None and s2 is not None:
        return _interpolate_global_objects_sample(s1, w1, s2, w2)
    return None
